#include "segmentationalgorithm.h"
using namespace std;
using namespace cv;
#define sqrtK 32		// 超像素个数32*32
#define sqrtN 512		// 图像格式512*512
int label[sqrtN][sqrtN];		// 图像各像素点归属
int dis[sqrtN][sqrtN];			// 图像各像素点距离
struct cluster {
    int row, col, l, a, b;
};
cluster clusters[sqrtK * sqrtK];		// 存储超像素的像素坐标、颜色
SegmentationAlgorithm::SegmentationAlgorithm() {}
/**
 * 初始化每一个超像素的坐标
 * 初始化每一个像素的label（即属于哪一个超像素）
 * 像素与超像素的距离先假设为-1
*/
void init_clusters(const Mat lab, int S) {
    for (int i = 0; i < sqrtK; i++) {
        int temp_row = S / 2 + i * S;
        for (int j = 0; j < sqrtK; j++) {
            clusters[i * sqrtK + j].row = temp_row;
            clusters[i * sqrtK + j].col = S / 2 + j * S;
        }
    }

    for (int i = 0; i < sqrtN; i++) {
        int cluster_row = i / S;
        for (int j = 0; j < sqrtN; j++) {
            label[i][j] = cluster_row * sqrtK + j / S;
        }
    }

    fill(dis[0], dis[0] + (sqrtN * sqrtN), -1);
}

inline int get_distance(const Mat lab, int clusters_index, int i, int j) {
    int dl = clusters[clusters_index].l - lab.at<Vec3b>(i, j)[0];
    int da = clusters[clusters_index].a - lab.at<Vec3b>(i, j)[1];
    int db = clusters[clusters_index].b - lab.at<Vec3b>(i, j)[2];
    int dx = clusters[clusters_index].row - i;
    int dy = clusters[clusters_index].col - j;

    int h_distance = dl * dl + da * da + db * db;
    int xy_distance = dx * dx + dy * dy;
    //cout << h_distance << "\t" << xy_distance * 100 << endl;
    return h_distance + xy_distance * 100;
}

void update_pixel(const Mat lab, int s) {
    for (int i = 0; i < sqrtK * sqrtK; i++) {	// 对于每一个超像素
        int clusters_x = clusters[i].row;
        int clusters_y = clusters[i].col;
        for (int x = -s; x <= s; x++) {			// 在它周围-s到s的范围内
            for (int y = -s; y <= s; y++) {
                int now_x = clusters_x + x;
                int now_y = clusters_y + y;
                if (now_x < 0 || now_x >= sqrtN || now_y < 0 || now_y >= sqrtN)
                    continue;
                int new_dis = get_distance(lab, i, now_x, now_y);
                // 如果为-1（还没有更新过）或者新的距离更小，就更换当前像素属于的超像素
                if (dis[now_x][now_y] > new_dis || dis[now_x][now_y] == -1) {
                    dis[now_x][now_y] = new_dis;
                    label[now_x][now_y] = i;
                }
            }
        }
    }
}

void updaye_clusters(const Mat lab) {
    int* sum_count = new int[sqrtK * sqrtK]();
    int* sum_i = new int[sqrtK * sqrtK]();
    int* sum_j = new int[sqrtK * sqrtK]();
    int* sum_l = new int[sqrtK * sqrtK]();
    int* sum_a = new int[sqrtK * sqrtK]();
    int* sum_b = new int[sqrtK * sqrtK]();
    for (int i = 0; i < sqrtN; i++) {
        for (int j = 0; j < sqrtN; j++) {
            sum_count[label[i][j]]++;
            sum_i[label[i][j]] += i;
            sum_j[label[i][j]] += j;
            sum_l[label[i][j]] += lab.at<Vec3b>(i, j)[0];
            sum_a[label[i][j]] += lab.at<Vec3b>(i, j)[1];
            sum_b[label[i][j]] += lab.at<Vec3b>(i, j)[2];
        }
    }
    for (int i = 0; i < sqrtK * sqrtK; i++) {
        if (sum_count[i] == 0) {
            continue;
        }
        clusters[i].row = round(sum_i[i] / sum_count[i]);
        clusters[i].col = round(sum_j[i] / sum_count[i]);
        clusters[i].l = round(sum_l[i] / sum_count[i]);
        clusters[i].a = round(sum_a[i] / sum_count[i]);
        clusters[i].b = round(sum_b[i] / sum_count[i]);
    }
    delete[] sum_count;
    delete[] sum_i;
    delete[] sum_j;
    delete[] sum_l;
    delete[] sum_a;
    delete[] sum_b;
}

void draw_clusters(const Mat copy) {
    for (int index = 0; index < sqrtK * sqrtK; index++) {
        Point p(clusters[index].row, clusters[index].col);
        circle(copy, p, 1, Scalar(0, 0, 255), 1);  // 画半径为1的圆(画点）
    }
    imshow("超像素示意图", copy);
}

void final_draw(const Mat lab, Mat copy) {
    for (int i = 0; i < sqrtN; i++) {
        for (int j = 0; j < sqrtN; j++) {
            int index = label[i][j];
            copy.at<Vec3b>(i, j)[0] = lab.at<Vec3b>(clusters[index].row, clusters[index].col)[0];
            copy.at<Vec3b>(i, j)[1] = lab.at<Vec3b>(clusters[index].row, clusters[index].col)[1];
            copy.at<Vec3b>(i, j)[2] = lab.at<Vec3b>(clusters[index].row, clusters[index].col)[2];
        }
    }
    cvtColor(copy, copy, COLOR_Lab2BGR);
    imshow("分割图", copy);
}

void draw_edge(const Mat lab, Mat copy) {
    // 这里的代码和上面的函数几乎一样，都是同标签的绘制相应的超像素颜色，因为方便用户自己选用绘制函数所以没有调用上面的函数
    for (int i = 0; i < sqrtN; i++) {
        for (int j = 0; j < sqrtN; j++) {
            int index = label[i][j];
            copy.at<Vec3b>(i, j)[0] = lab.at<Vec3b>(clusters[index].row, clusters[index].col)[0];
            copy.at<Vec3b>(i, j)[1] = lab.at<Vec3b>(clusters[index].row, clusters[index].col)[1];
            copy.at<Vec3b>(i, j)[2] = lab.at<Vec3b>(clusters[index].row, clusters[index].col)[2];
        }
    }

    // 这里的思路是4个方向，一旦有标签不同，就设置为黑色
    static int X[] = { 0,0,-1,1 };
    static int Y[] = { 1,-1,0,0 };
    cvtColor(copy, copy, COLOR_Lab2BGR); // 改成BGR，方便后面设置边框的颜色。
    for (int i = 0; i < sqrtN; i++) {
        for (int j = 0; j < sqrtN; j++) {
            int index = label[i][j];
            for (int k = 0; k < 4; k++) {
                if (index != label[i + X[k]][j + X[k]])
                    copy.at<Vec3b>(i, j)[0] = copy.at<Vec3b>(i, j)[1] = copy.at<Vec3b>(i, j)[2] = 0;
            }
        }
    }
    imshow("超像素边界", copy);
}

int aaaa() {
    // 注意修改文件位置
    Mat src = imread("C:\\Users\\ttp\\Desktop\\map.bmp"), lab;

    // resize图片并高斯滤波（可选）
    resize(src, src, Size(sqrtN, sqrtN));
    // GaussianBlur(src, src, Size(3, 3), 1, 1);

    /**
        得到Lab色彩空间，需要注意的是:
        1.opencv里面默认为BGR排列方式
        2.LAB通道范围取决于转换前的通道范围，这样其实也方便处理
        例如：开始是0-255，转换后也是0-255，而不是LAB规定的[127,-128]
    */
    // cvtColor(src, lab, cv::CV_BGR2Lab);
    cvtColor(src, lab, COLOR_BGR2Lab);
    int N = sqrtN * sqrtN;			// 像素总数 512*512
    int K = sqrtK * sqrtK;			// 超像素个数 128*128
    int S = sqrt(N / K);			// 相邻种子点距离（超像素边长） 4

    // 1.初始化像素
    init_clusters(lab, S);
    cout << "1-初始化像素-完成\n";

    // 2.微调种子的位置 貌似好一点,没有太大区别，所以这里就直接省略了

    for (int i = 0; i < 5; i++) {
        // 3.4.初始化数据
        update_pixel(lab, 2 * S);
        cout << "3-初始化数据-完成\n";

        // 5.让超像素位于正中间
        updaye_clusters(lab);
        cout << "4-让超像素位于正中间-完成\n";

        // -------------------这两个函数主要是帮助显示结果的
        // 6.标识超像素
        draw_clusters(src.clone());
        cout << "5-标识超像素-完成\n";

        // 7.绘制超像素结果图
        final_draw(lab, lab.clone());
        cout << "6-绘制超像素结果图-完成\n";

        draw_edge(lab, lab.clone());

        // opencv的函数，每1000ms更新一下，动态显示图片
        waitKey(30);
        // -----------------------------------------------
    }
    imshow("原图", src);
    waitKey(0);
}
